import numpy as np
import torch
import matplotlib.pyplot as plt
from sklearn.metrics import precision_recall_curve, auc
import models.white_label_model as models
import utils.datasets_1d as data_utils
from torchvision import transforms
from torch.utils.data import DataLoader
from tqdm import tqdm

model_old_path = (
    "/home/ubuntu/cxf/keypoint_detection/runs/train_cxf_5/sigma_2.05epoch_15/model.pt"
)


model_new_path = "/home/ubuntu/cxf/keypoint_detection/bhw_result/02/best_model.pt"
validation_dataset_dir = "/home/ubuntu/cxf/keypoint_detection/bhw_dataset/val"
input_size: tuple = (512, 384)
output_size: tuple = (512, 384)
heatmap_sigma = 3
num_workers = 50
device = torch.device("cuda:2" if torch.cuda.is_available() else "cpu")
fig_save_path = "/home/ubuntu/cxf/keypoint_detection/bhw_result/PR_curve.png"


# model_old = models.GrayDemoNet2().to(device)
model_old = torch.load(model_old_path).to(device)
# model1.load_state_dict(checkpoint['model_state_dict'])

model_new = models.GrayDemoNet2().to(device)
checkpoint = torch.load(model_new_path)
model_new.load_state_dict(checkpoint["model_state_dict"])


transform = transforms.ToTensor()

val_dataset = data_utils.WhiteLabelKeypointDataset(
    validation_dataset_dir,
    transform=transform,
    input_size=input_size,
    output_size=output_size,
    sigma=heatmap_sigma,
)
val_loader = DataLoader(
    val_dataset,
    batch_size=1,
    shuffle=False,
    collate_fn=data_utils.collate_fn,
    num_workers=num_workers,
)


def process_output(output: torch.Tensor, heatmap: torch.Tensor):
    pred_heatmap = output[:, 0, :, :]
    true_heatmap = heatmap[:, 0, :, :]
    true_heatmap[true_heatmap > 0.5] = 1
    true_heatmap[true_heatmap <= 0.5] = 0

    return true_heatmap, pred_heatmap


total_true_heatmap_old = torch.zeros(1).to(device)
total_pred_heatmap_old = torch.zeros(1).to(device)
total_true_heatmap_new = torch.zeros(1).to(device)
total_pred_heatmap_new = torch.zeros(1).to(device)

for image, heatmap in tqdm(val_loader):
    image = image.to(device)
    heatmap = heatmap.to(device)
    with torch.no_grad():
        output_old = model_old(image)
        output_new = model_new(image)

        true_heatmap_old, pred_heatmap_old = process_output(output_old, heatmap)
        true_heatmap_new, pred_heatmap_new = process_output(output_new, heatmap)
        total_true_heatmap_old = torch.cat(
            (total_true_heatmap_old, true_heatmap_old.flatten())
        )
        total_pred_heatmap_old = torch.cat(
            (total_pred_heatmap_old, pred_heatmap_old.flatten())
        )
        total_true_heatmap_new = torch.cat(
            (total_true_heatmap_new, true_heatmap_new.flatten())
        )
        total_pred_heatmap_new = torch.cat(
            (total_pred_heatmap_new, pred_heatmap_new.flatten())
        )

total_true_heatmap_old = total_true_heatmap_old.cpu().numpy()
total_pred_heatmap_old = total_pred_heatmap_old.cpu().numpy()
total_true_heatmap_new = total_true_heatmap_new.cpu().numpy()
total_pred_heatmap_new = total_pred_heatmap_new.cpu().numpy()


assert (
    len(total_true_heatmap_old)
    == len(total_pred_heatmap_old)
    == len(total_true_heatmap_new)
    == len(total_pred_heatmap_new)
)


precision_old, recall_old, thresholds_old = precision_recall_curve(
    total_true_heatmap_old[1:], total_pred_heatmap_old[1:]
)
pr_auc_old = auc(recall_old, precision_old)
precision_new, recall_new, thresholds_new = precision_recall_curve(
    total_true_heatmap_new[1:], total_pred_heatmap_new[1:]
)
pr_auc_new = auc(recall_new, precision_new)


for i in tqdm(range(len(precision_old))):
    if abs(precision_old[i] - recall_old[i]) < 1e-16:
        idx_old = i

for i in tqdm(range(len(precision_new))):
    if abs(precision_new[i] - recall_new[i]) < 1e-16:
        idx_new = i

print(f"idx_old: {idx_old}, idx_new: {idx_new}")


# 绘制PR曲线
plt.figure(figsize=(6, 6))
plt.plot(
    recall_old, precision_old, color="r", label=f"PR Curve Old (AUC = {pr_auc_old:.3f})"
)
plt.plot(
    recall_new, precision_new, color="b", label=f"PR Curve New (AUC = {pr_auc_new:.3f})"
)
plt.plot([0, 1], [0, 1], color="gray", linestyle="--")
plt.scatter(
    recall_old[idx_old],
    precision_old[idx_old],
    marker="o",
    color="r",
    label=f"Threshold = {thresholds_old[idx_old]:.3f}",
)
plt.scatter(
    recall_new[idx_new],
    precision_new[idx_new],
    marker="o",
    color="b",
    label=f"Threshold = {thresholds_new[idx_new]:.3f}",
)
plt.xlabel("Recall")
plt.ylabel("Precision")
plt.title("Precision-Recall Curve")
plt.legend()
plt.grid()
plt.savefig(fig_save_path)
plt.show()
